import os
import re

IGNORE_DIRS = {"prompt"}          # 可在这里继续加入要忽略的目录名（小写）

def natural_keys(text: str):
    """把字符串拆分成字符串+数字列表，用于自然排序（DM01A < DM01B）"""
    return [int(s) if s.isdigit() else s.lower()
            for s in re.split(r'(\d+)', text)]

def find_md_files_sorted(base_dir='.'):
    """递归查找 .md，README 优先，同级自然排序，并忽略指定目录"""
    all_files = []
    for root, dirs, files in os.walk(base_dir):
        # ── 1️⃣ 过滤掉需要忽略的子目录 ────────────────────────────────
        dirs[:] = [d for d in dirs if d.lower() not in IGNORE_DIRS]

        # 2️⃣ 让 os.walk 的子目录遍历顺序稳定（自然排序）
        dirs.sort(key=natural_keys)

        # 3️⃣ 取当前目录下的 .md 文件
        md_files = [f for f in files if f.endswith('.md')]

        # 4️⃣ README.md 优先，其余按自然顺序
        md_files.sort(key=lambda x: (x.lower() != 'readme.md', natural_keys(x)))

        for f in md_files:
            rel_path = os.path.relpath(os.path.join(root, f), base_dir).replace('\\', '/')

            # 5️⃣ 双保险：若路径里包含被忽略目录，也跳过
            if any(f"/{ign}/" in rel_path.lower() for ign in IGNORE_DIRS):
                continue

            all_files.append(rel_path)

    # 6️⃣ 全局再做一次自然排序，保证最终稳定
    all_files.sort(key=natural_keys)
    return all_files

def extract_title(md_file_path):
    """提取 markdown 文件中首个 # 标题；若无标题，用文件名"""
    try:
        with open(md_file_path, 'r', encoding='utf-8') as f:
            for line in f:
                if line.strip().startswith('#'):
                    return line.lstrip('#').strip()
    except Exception:
        pass
    return os.path.basename(md_file_path)

def generate_toc(md_files):
    """生成 markdown TOC，缩进表示目录层级"""
    toc_lines = ["## 📚 目录", ""]
    for path in md_files:
        if path.lower() == 'readme.md':
            continue                               # 跳过仓库根 README.md
        depth = path.count('/')
        indent = '  ' * depth
        title = extract_title(path)
        link = path.replace(' ', '%20')            # URL 编码空格
        toc_lines.append(f"{indent}- [{title}]({link})")
    toc_lines.append("")                           # 结尾空行
    return "\n".join(toc_lines)

def insert_toc_to_readme(toc_md, readme_path='README.md'):
    """把 TOC 插入 README.md 的 <!-- TOC START -->…<!-- TOC END --> 区域"""
    if not os.path.exists(readme_path):
        with open(readme_path, 'w', encoding='utf-8') as f:
            f.write("# 项目说明\n\n")

    with open(readme_path, 'r', encoding='utf-8') as f:
        content = f.read()

    start_tag, end_tag = "<!-- TOC START -->", "<!-- TOC END -->"
    if start_tag in content and end_tag in content:
        pre = content.split(start_tag)[0]
        post = content.split(end_tag)[1]
    else:
        pre, post = content.rstrip() + "\n\n", ""

    new_content = f"{pre}{start_tag}\n{toc_md}\n{end_tag}\n{post}"
    with open(readme_path, 'w', encoding='utf-8') as f:
        f.write(new_content)

if __name__ == '__main__':
    md_files = find_md_files_sorted()
    toc_markdown = generate_toc(md_files)
    insert_toc_to_readme(toc_markdown)
    print("✅ 生成完成：TOC 已插入 README.md（README 优先，同级自然排序，忽略 PROMPT 目录）")